# viz/sankey.py —— 字段血缘（表 -> 列） + 列到列桑基
from typing import Dict, Iterable, Tuple, List, Set, Optional
from collections import defaultdict

from pyecharts import options as opts
from pyecharts.charts import Sankey
from core.models import StatementLineage


def _rename_star_group(cols_by_tbl: Dict[str, Iterable[str]], derived_name: str) -> Dict[str, List[str]]:
    out: Dict[str, List[str]] = {}
    for k, v in cols_by_tbl.items():
        key = derived_name if k == "*" else k
        out[key] = sorted(list(v))
    return out


def _limit_columns_per_table(cols_by_tbl: Dict[str, List[str]], top_n: Optional[int]) -> Dict[str, List[str]]:
    if not top_n or top_n <= 0:
        return cols_by_tbl
    limited: Dict[str, List[str]] = {}
    for tbl, cols in cols_by_tbl.items():
        if len(cols) <= top_n:
            limited[tbl] = cols
        else:
            head, tail = cols[:top_n], cols[top_n:]
            limited[tbl] = head + (["(others)"] if tail else [])
    return limited


def _node_color(name: str) -> Optional[str]:
    n = name.lower()
    if n.startswith("derived"):
        return "#ffb86b"   # 橙
    if n.startswith("hive."):
        return "#8be9fd"   # 青
    if n.startswith("mysql"):
        return "#bd93f9"   # 紫
    if n.startswith("subquery:"):
        return "#50fa7b"   # 绿
    return None


def _build_nodes_links(cols_by_tbl: Dict[str, List[str]]) -> Tuple[List[dict], List[dict]]:
    nodes: List[dict] = []
    links: List[dict] = []
    seen_nodes: Set[str] = set()
    edge_weight: Dict[Tuple[str, str], int] = defaultdict(int)

    for tbl, cols in sorted(cols_by_tbl.items()):
        if tbl not in seen_nodes:
            color = _node_color(tbl)
            node = {"name": tbl}
            if color:
                node["itemStyle"] = {"color": color}
            nodes.append(node)
            seen_nodes.add(tbl)
        for c in cols:
            name = f"{tbl}.{c}"
            if name not in seen_nodes:
                color = _node_color(tbl)  # 列节点沿用表色
                node = {"name": name}
                if color:
                    node["itemStyle"] = {"color": color}
                nodes.append(node)
                seen_nodes.add(name)
            edge_weight[(tbl, name)] += 1

    for (src, dst), w in edge_weight.items():
        links.append({"source": src, "target": dst, "value": w})

    return nodes, links


def render_column_sankey(
    st: StatementLineage,
    *,
    title: str = "字段血缘（表 → 列）",
    derived_group_name: str = "Derived",
    top_n_columns: int = 20,
    width: str = "1200px",
    height: str = "700px",
    return_html: bool = True,
):
    chart = Sankey(init_opts=opts.InitOpts(width=width, height=height))
    if not st.columns_by_table:
        chart.set_global_opts(title_opts=opts.TitleOpts(title=f"{title}（无字段）"))
        return chart.render_embed() if return_html else chart

    renamed = _rename_star_group(st.columns_by_table, derived_group_name)
    limited = _limit_columns_per_table(renamed, top_n_columns)
    nodes, links = _build_nodes_links(limited)

    chart = (
        Sankey(init_opts=opts.InitOpts(width=width, height=height))
        .add(
            "表与字段",
            nodes=nodes,
            links=links,
            linestyle_opt=opts.LineStyleOpts(opacity=0.45, curve=0.5, color="source"),
            label_opts=opts.LabelOpts(position="right", formatter="{b}"),
            node_width=16,
            node_gap=10,
            tooltip_opts=opts.TooltipOpts(formatter="{b}"),
        )
        .set_global_opts(title_opts=opts.TitleOpts(title=title))
    )
    return chart.render_embed() if return_html else chart


# ✨ 新增：列到列（源列 → 目标列名/别名）桑基
def render_column_to_column_sankey(
    st: StatementLineage,
    *,
    title: str = "字段血缘（源列 → 目标列）",
    width: str = "1200px",
    height: str = "700px",
    return_html: bool = True,
):
    lineage = st.column_lineage or {}
    chart = Sankey(init_opts=opts.InitOpts(width=width, height=height))
    if not lineage:
        chart.set_global_opts(title_opts=opts.TitleOpts(title=f"{title}（无数据）"))
        return chart.render_embed() if return_html else chart

    nodes: List[dict] = []
    links: List[dict] = []
    seen: Set[str] = set()
    weights: Dict[Tuple[str, str], int] = defaultdict(int)

    for src, tgts in lineage.items():
        if src not in seen:
            color = _node_color(src.split(".", 1)[0])  # 以“表前缀”着色
            node = {"name": src}
            if color:
                node["itemStyle"] = {"color": color}
            nodes.append(node)
            seen.add(src)
        for tgt in sorted(tgts):
            if tgt not in seen:
                # 目标列（别名）作为裸节点
                nodes.append({"name": tgt})
                seen.add(tgt)
            weights[(src, tgt)] += 1

    for (sname, tname), w in weights.items():
        links.append({"source": sname, "target": tname, "value": w})

    chart = (
        Sankey(init_opts=opts.InitOpts(width=width, height=height))
        .add(
            "列到列",
            nodes=nodes,
            links=links,
            linestyle_opt=opts.LineStyleOpts(opacity=0.45, curve=0.5, color="source"),
            label_opts=opts.LabelOpts(position="right", formatter="{b}"),
            node_width=16,
            node_gap=10,
            tooltip_opts=opts.TooltipOpts(formatter="{b}"),
        )
        .set_global_opts(title_opts=opts.TitleOpts(title=title))
    )
    return chart.render_embed() if return_html else chart
